/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to you under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.calcite.rel.logical;

import com.google.common.collect.Iterables;
import com.google.common.collect.LinkedListMultimap;
import com.google.common.collect.Lists;
import com.google.common.collect.Multimap;
import org.apache.calcite.linq4j.Ord;
import org.apache.calcite.plan.RelOptCluster;
import org.apache.calcite.plan.RelOptCost;
import org.apache.calcite.plan.RelOptPlanner;
import org.apache.calcite.plan.RelOptUtil;
import org.apache.calcite.plan.RelTraitSet;
import org.apache.calcite.rel.RelCollation;
import org.apache.calcite.rel.RelNode;
import org.apache.calcite.rel.RelWriter;
import org.apache.calcite.rel.core.Window;
import org.apache.calcite.rel.externalize.RelDrdsWriter;
import org.apache.calcite.rel.externalize.RexExplainVisitor;
import org.apache.calcite.rel.metadata.RelMetadataQuery;
import org.apache.calcite.rel.type.RelDataType;
import org.apache.calcite.rel.type.RelDataTypeField;
import org.apache.calcite.rex.RexCall;
import org.apache.calcite.rex.RexFieldCollation;
import org.apache.calcite.rex.RexInputRef;
import org.apache.calcite.rex.RexLiteral;
import org.apache.calcite.rex.RexLocalRef;
import org.apache.calcite.rex.RexNode;
import org.apache.calcite.rex.RexOver;
import org.apache.calcite.rex.RexProgram;
import org.apache.calcite.rex.RexShuttle;
import org.apache.calcite.rex.RexWindow;
import org.apache.calcite.rex.RexWindowBound;
import org.apache.calcite.runtime.PredicateImpl;
import org.apache.calcite.sql.SqlAggFunction;
import org.apache.calcite.tools.RelBuilder;
import org.apache.calcite.util.ImmutableBitSet;
import org.apache.calcite.util.Litmus;
import org.apache.calcite.util.Pair;
import org.apache.commons.lang.StringUtils;

import java.util.AbstractList;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.IdentityHashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * Sub-class of {@link org.apache.calcite.rel.core.Window}
 * not targeted at any particular engine or calling convention.
 */
public final class LogicalWindow extends Window {
  /**
   * Creates a LogicalWindow.
   *
   * <p>Use {@link #create} unless you know what you're doing.
   *
   * @param cluster Cluster
   * @param traitSet Trait set
   * @param input   Input relational expression
   * @param constants List of constants that are additional inputs
   * @param rowType Output row type
   * @param groups Window groups
   */
  public LogicalWindow(RelOptCluster cluster, RelTraitSet traitSet,
      RelNode input, List<RexLiteral> constants, RelDataType rowType,
      List<Group> groups) {
    super(cluster, traitSet, input, constants, rowType, groups);
  }

  @Override public LogicalWindow copy(RelTraitSet traitSet,
      List<RelNode> inputs) {
    LogicalWindow logicalWindow =  new LogicalWindow(getCluster(), traitSet, sole(inputs), constants,
      rowType, groups);
    logicalWindow.setFixedCost(getFixedCost());
    return logicalWindow;
  }

  /**
   * Creates a LogicalWindow.
   *
   * @param input   Input relational expression
   * @param traitSet Trait set
   * @param constants List of constants that are additional inputs
   * @param rowType Output row type
   * @param groups Window groups
   */
  public static LogicalWindow create(RelTraitSet traitSet, RelNode input,
      List<RexLiteral> constants, RelDataType rowType, List<Group> groups) {
    return new LogicalWindow(input.getCluster(), traitSet, input, constants,
        rowType, groups);
  }

  @Override public RelOptCost computeSelfCost(RelOptPlanner planner, RelMetadataQuery mq) {
    if (getFixedCost() != null) {
      return getFixedCost();
    }
    // Cost is proportional to the number of rows and the number of
    // components (groups and aggregate functions). There is
    // no I/O cost.
    //
    // TODO #1. Add memory cost.
    // TODO #2. MIN and MAX have higher CPU cost than SUM and COUNT.
    final double rowsIn = mq.getRowCount(getInput());
    int count = groups.size();
    for (Group group : groups) {
      count += group.aggCalls.size();
    }
    return planner.getCostFactory().makeCost(rowsIn, rowsIn * count, 0, 0, 0);
  }

//  public RelOptCost computeSelfCost(RelOptPlanner planner) {
//    return computeSelfCost(planner, RelMetadataQuery.instance());
//  }

  /**
   * Creates a LogicalWindow by parsing a {@link RexProgram}.
   */
  public static RelNode create(RelOptCluster cluster,
      RelTraitSet traitSet, RelBuilder relBuilder, RelNode child,
      final RexProgram program) {
    final RelDataType outRowType = program.getOutputRowType();
    // Build a list of distinct groups, partitions and aggregate
    // functions.
    final Multimap<WindowKey, RexOver> windowMap = LinkedListMultimap.create();

    final int inputFieldCount = child.getRowType().getFieldCount();

    final Map<RexLiteral, RexInputRef> constantPool = new HashMap<>();
    final List<RexLiteral> constants = new ArrayList<>();

    // Identify constants in the expression tree and replace them with
    // references to newly generated constant pool.
      // For: RexSlot,RexProgram,Over index
    RexShuttle replaceConstants = new RexShuttle() {
      @Override public RexNode visitLiteral(RexLiteral literal) {
        RexInputRef ref = constantPool.get(literal);
        if (ref != null) {
          return ref;
        }
        constants.add(literal);
        ref = new RexInputRef(constantPool.size() + inputFieldCount,
            literal.getType());
        constantPool.put(literal, ref);
        return ref;
      }
    };

    class RexShuttleCorrectConstants extends RexShuttle {
      final private List<RexLiteral> correctedConstants = new ArrayList<>();
      final Map<RexInputRef, RexInputRef> correctedConstantPool = new HashMap<>();
      final private int inputFieldCount;
      final private int adjustIndex;
      public RexShuttleCorrectConstants(int inputFieldCount, int adjustIndex) {
        this.inputFieldCount = inputFieldCount;
        this.adjustIndex = adjustIndex;
      }
      
      @Override public RexNode visitInputRef(RexInputRef inputRef) {
        if (!constantPool.containsValue(inputRef)) {
          return inputRef;
        }
        if (correctedConstantPool.containsKey(inputRef)) {
          return correctedConstantPool.get(inputRef);
        }
        int index = inputRef.getIndex() - inputFieldCount;
        RexLiteral rexLiteral = constants.get(index);
        RexInputRef ref = new RexInputRef(correctedConstants.size() + inputFieldCount + adjustIndex,
            rexLiteral.getType());
        correctedConstants.add(rexLiteral);
        correctedConstantPool.put(inputRef, ref);
        return ref;
      }
    }
    
    // Build a list of groups, partitions, and aggregate functions. Each
    // aggregate function will add its arguments as outputs of the input
    // program.
    final Map<RexOver, RexOver> origToNewOver = new IdentityHashMap<>();
    for (RexNode agg : program.getExprList()) {
      if (agg instanceof RexOver) {
        final RexOver origOver = (RexOver) agg;
        final RexOver newOver = (RexOver) origOver.accept(replaceConstants);
        origToNewOver.put(origOver, newOver);
        addWindows(windowMap, newOver, inputFieldCount);
      }
    }

    final Map<RexOver, Window.RexWinAggCall> aggMap = new HashMap<>();
    List<Group> groups = new ArrayList<>();
    List<WindowKey> listKeys = Arrays.asList( windowMap.keySet().toArray(new WindowKey[0]));
    Collections.sort(listKeys, new Comparator<WindowKey>() {
      @Override
      public int compare(WindowKey o1, WindowKey o2) {

          WindowSortedKey sortedKey1 = new WindowSortedKey(o1);
          WindowSortedKey sortedKey2 = new WindowSortedKey(o2);
        int i = sortedKey1.compareTo(sortedKey2);
        return i;
      }
    });
    int[] groupSplit = new int[listKeys.size()];
    int groupPosition = 0;
    WindowKey preWindowKey = null;
    for (WindowKey windowKey : listKeys) {
      if (preWindowKey == null) {
        preWindowKey = windowKey;
        ++ groupSplit[groupPosition];
      } else {
        if (preWindowKey.equals(windowKey)) {
          ++ groupSplit[groupPosition];
        } else {
          ++ groupSplit[++groupPosition];
          preWindowKey = windowKey;
        }
      }
      final List<RexWinAggCall> aggCalls = new ArrayList<>();
      for (RexOver over : windowMap.get(windowKey)) {
        final RexWinAggCall aggCall =
            new RexWinAggCall(
                over.getAggOperator(),
                over.getType(),
                toInputRefs(over.operands),
                aggMap.size(),
                over.isDistinct());
        aggCalls.add(aggCall);
        aggMap.put(over, aggCall);
      }
      RexShuttle toInputRefs = new RexShuttle() {
        @Override public RexNode visitLocalRef(RexLocalRef localRef) {
          return new RexInputRef(localRef.getIndex(), localRef.getType());
        }
      };
      groups.add(
          new Group(
              windowKey.groupSet,
              windowKey.isRows,
              windowKey.lowerBound.accept(toInputRefs),
              windowKey.upperBound.accept(toInputRefs),
              windowKey.orderKeys,
              aggCalls));
    }
    
    // Figure out the type of the inputs to the output program.
    // They are: the inputs to this rel, followed by the outputs of
    // each window.
    final List<Window.RexWinAggCall> flattenedAggCallList = new ArrayList<>();
    final List<Map.Entry<String, RelDataType>> fieldList =
        new ArrayList<Map.Entry<String, RelDataType>>(
            child.getRowType().getFieldList());
    final int offset = fieldList.size();

    // Use better field names for agg calls that are projected.
    final Map<Integer, String> fieldNames = new HashMap<>();
    for (Ord<RexLocalRef> ref : Ord.zip(program.getProjectList())) {
      final int index = ref.e.getIndex();
      if (index >= offset) {
        fieldNames.put(
            index - offset, outRowType.getFieldNames().get(ref.i));
      }
    }

    for (Ord<Group> window : Ord.zip(groups)) {
      for (Ord<RexWinAggCall> over : Ord.zip(window.e.aggCalls)) {
        // Add the k-th over expression of
        // the i-th window to the output of the program.
        String name = fieldNames.get(over.i);
        if (name == null || name.startsWith("$")) {
          name = "f" + fieldList.size() + "w" + window.i + "$o" + over.i;
        }
        fieldList.add(Pair.of(name, over.e.getType()));
        flattenedAggCallList.add(over.e);
      }
    }
    final RelDataType intermediateRowType =
        cluster.getTypeFactory().createStructType(fieldList);

    // The output program is the windowed agg's program, combined with
    // the output calc (if it exists).
    RexShuttle shuttle =
        new RexShuttle() {
          @Override
          public RexNode visitOver(RexOver over) {
            // Look up the aggCall which this expr was translated to.
            final Window.RexWinAggCall aggCall =
                aggMap.get(origToNewOver.get(over));
            assert aggCall != null;
            assert RelOptUtil.eq(
                "over",
                over.getType(),
                "aggCall",
                aggCall.getType(),
                Litmus.THROW);

            // Find the index of the aggCall among all partitions of all
            // groups.
            final int aggCallIndex =
                flattenedAggCallList.indexOf(aggCall);
            assert aggCallIndex >= 0;

            // Replace expression with a reference to the window slot.
            final int index = inputFieldCount + aggCallIndex;
            assert RelOptUtil.eq(
                "over",
                over.getType(),
                "intermed",
                intermediateRowType.getFieldList().get(index).getType(),
                Litmus.THROW);
            return new RexInputRef(
                index,
                over.getType());
          }

          @Override
          public RexNode visitLocalRef(RexLocalRef localRef) {
            final int index = localRef.getIndex();
            if (index < inputFieldCount) {
              // Reference to input field.
              return localRef;
            }
            return new RexLocalRef(
                flattenedAggCallList.size() + index,
                localRef.getType());
          }
        };
    RelNode childRel = child;
    int from = 0,to = 0;
    for (int i = 0; i < groupSplit.length; i++) {
      if (i >= 1) {
        from = to;
      }
      to += groupSplit[i];
      int childInputFieldCount = childRel.getRowType().getFieldCount();
      final List<Group> tempGroups = groups.subList(from,to);
      final List<Group> newGroups = new ArrayList<>();
      final AtomicInteger aggSize = new AtomicInteger();
      AtomicInteger aggCallIndex = new AtomicInteger();
      final RexShuttleCorrectConstants correctConstants = new RexShuttleCorrectConstants(inputFieldCount, childRel.getRowType().getFieldCount() - inputFieldCount);
      tempGroups.stream().forEach(t -> {
        RexWindowBound lowerBound = t.lowerBound.accept(correctConstants);
        RexWindowBound upperBound = t.upperBound.accept(correctConstants);
        aggSize.set(aggSize.get() + t.aggCalls.size());
        final List<RexWinAggCall> newAggCalls = new ArrayList<>();
        for (RexWinAggCall tempAggCall:t.aggCalls) {
          //reajust the index
          RexCall newRexCall = (RexCall) tempAggCall.accept(correctConstants);
          final RexWinAggCall newAggCall =
              new RexWinAggCall(
                  (SqlAggFunction) newRexCall.op,
                  newRexCall.type,
                  toInputRefs(newRexCall.operands),
                  aggCallIndex.getAndIncrement(),
                  tempAggCall.distinct);
          newAggCalls.add(newAggCall);
        }
        newGroups.add(new Group(
            t.keys,
            t.isRows,
            lowerBound,
            upperBound,
            t.orderKeys,
            newAggCalls));
      });
      final RelDataType tempIntermediateRowType =
          cluster.getTypeFactory().createStructType(intermediateRowType.getFieldList().subList(0, childInputFieldCount + aggSize.get()));
      final LogicalWindow window =
          LogicalWindow.create(traitSet, childRel, correctConstants.correctedConstants, tempIntermediateRowType,
              newGroups);
      childRel = window;

    }

    // The order that the "over" calls occur in the groups and
    // partitions may not match the order in which they occurred in the
    // original expression.
    // Add a project to permute them.
    final List<RexNode> rexNodesWindow = new ArrayList<>();
    for (RexNode rexNode : program.getExprList()) {
      rexNodesWindow.add(rexNode.accept(shuttle));
    }
    final List<RexNode> refToWindow = toInputRefs(rexNodesWindow);

    final List<RexNode> projectList = new ArrayList<>();
    for (RexLocalRef inputRef : program.getProjectList()) {
      final int index = inputRef.getIndex();
      final RexInputRef ref = (RexInputRef) refToWindow.get(index);
      projectList.add(ref);
    }

    return relBuilder.push(childRel)
        .project(projectList, outRowType.getFieldNames())
        .build();
  }

  private static List<RexNode> toInputRefs(
      final List<? extends RexNode> operands) {
    return new AbstractList<RexNode>() {
      public int size() {
        return operands.size();
      }

      public RexNode get(int index) {
        final RexNode operand = operands.get(index);
        if (operand instanceof RexInputRef) {
          return operand;
        }
        assert operand instanceof RexLocalRef;
        final RexLocalRef ref = (RexLocalRef) operand;
        return new RexInputRef(ref.getIndex(), ref.getType());
      }
    };
  }

  /** Group specification. All windowed aggregates over the same window
   * (regardless of how it is specified, in terms of a named window or specified
   * attribute by attribute) will end up with the same window key. */
  private static class WindowKey {
    private final ImmutableBitSet groupSet;
    private final RelCollation orderKeys;
    private final boolean isRows;
    private final RexWindowBound lowerBound;
    private final RexWindowBound upperBound;

    WindowKey(
        ImmutableBitSet groupSet,
        RelCollation orderKeys,
        boolean isRows,
        RexWindowBound lowerBound,
        RexWindowBound upperBound) {
      this.groupSet = groupSet;
      this.orderKeys = orderKeys;
      this.isRows = isRows;
      this.lowerBound = lowerBound;
      this.upperBound = upperBound;
    }

    @Override public int hashCode() {
      return Objects.hash(groupSet, orderKeys, isRows, lowerBound, upperBound);
    }

    @Override public boolean equals(Object obj) {
      return obj == this
          || obj instanceof WindowKey
          && groupSet.equals(((WindowKey) obj).groupSet)
          && orderKeys.equals(((WindowKey) obj).orderKeys)
          && Objects.equals(lowerBound, ((WindowKey) obj).lowerBound)
          && Objects.equals(upperBound, ((WindowKey) obj).upperBound)
          && isRows == ((WindowKey) obj).isRows;
    }
  }

  /** sorted the window key. */
  private static class WindowSortedKey implements Comparable{
    private final WindowKey windowKey;

    WindowSortedKey(WindowKey windowKey) {
      this.windowKey = windowKey;
    }

    @Override public int hashCode() {
      return Objects.hash(this.windowKey.groupSet, this.windowKey.orderKeys);
    }

    @Override public boolean equals(Object obj) {
      return obj == this
          || obj instanceof WindowKey
          && this.windowKey.groupSet.equals(((WindowSortedKey) obj).windowKey.groupSet)
          && this.windowKey.orderKeys.equals(((WindowSortedKey) obj).windowKey.orderKeys);
    }

    @Override
    public int compareTo(Object o) {
      return this.equals(o) ? -1 : 0;
    }
  }
  

  private static void addWindows(
      Multimap<WindowKey, RexOver> windowMap,
      RexOver over, final int inputFieldCount) {
    final RexWindow aggWindow = over.getWindow();

    // Look up or create a window.
    RelCollation orderKeys = getCollation(
        Lists.newArrayList(
            Iterables.filter(aggWindow.orderKeys,
              new PredicateImpl<RexFieldCollation>() {
                public boolean test(RexFieldCollation rexFieldCollation) {
                  // If ORDER BY references constant (i.e. RexInputRef),
                  // then we can ignore such ORDER BY key.
                  return rexFieldCollation.left instanceof RexLocalRef;
                }
              })));
    ImmutableBitSet groupSet =
        ImmutableBitSet.of(getProjectOrdinals(aggWindow.partitionKeys));
    final int groupLength = groupSet.length();
    if (inputFieldCount < groupLength) {
      // If PARTITION BY references constant, we can ignore such partition key.
      // All the inputs after inputFieldCount are literals, thus we can clear.
      groupSet =
          groupSet.except(ImmutableBitSet.range(inputFieldCount, groupLength));
    }

    WindowKey windowKey =
        new WindowKey(
            groupSet, orderKeys, aggWindow.isRows(),
            aggWindow.getLowerBound(), aggWindow.getUpperBound());
    windowMap.put(windowKey, over);
  }
  
  @Override
  public RelWriter explainTermsForDisplay(RelWriter pw) {
    pw.item(RelDrdsWriter.REL_NAME, "Window");
    pw.item("field num", getInput().getRowType().getFieldCount());
    for (Ord<RelDataTypeField> field : Ord.zip(getRowType().getFieldList())) {
      String fieldName = field.e.getName();
      if (fieldName == null) {
        fieldName = "field#" + field.i;
      }
    }
    
    for (Group group : groups) {
      RexExplainVisitor visitor = new RexExplainVisitor(this);
      List<String> groupList = new ArrayList<String>(this.groups.get(0).keys.length());
      for (int groupIndex : group.keys.asList()) {
        groupList.add(visitor.getField(groupIndex).getKey());
      }
      pw.itemIf("group", StringUtils.join(groupList, ",") + "=" + group.keys.asList(), !groupList.isEmpty());
      for (RexWinAggCall aggCall : group.aggCalls) {
        aggCall.accept(visitor);
        pw.item(rowType.getFieldList().get(0).getKey(), visitor.toSqlString());
      }
      pw.itemIf("orders", group.orderKeys.toString(),
          group.orderKeys.getFieldCollations() != null && group.orderKeys.getFieldCollations().size() > 0);
      for (RexWinAggCall aggCall : group.aggCalls) {
        aggCall.accept(visitor);
        pw.item(rowType.getFieldList().get(0).getKey(), visitor.toSqlString());
      }

      pw.itemIf("constants", constants.toString(), constants != null && constants.size() > 0);
      for (RexWinAggCall aggCall : group.aggCalls) {
        aggCall.accept(visitor);
        pw.item(rowType.getFieldList().get(0).getKey(), visitor.toSqlString());
      }
    }
    return pw;
  }
  
}

// End LogicalWindow.java
